import streamlit as st
import base64
from zai import ZhipuAiClient
from dotenv import load_dotenv
load_dotenv()

def llm_chat(message_list):
    client = ZhipuAiClient()
    msgs = []
    for message in message_list:
        content = message[1]
        role = message[0]
        img_base = None
        if len(message)>2:
            img_base = message[2]
        if role == 'ai':
            msgs.append({'role':'assistant','content':content})
        elif role == 'human':
            content_list = [
                {
                    "type": "text",
                    "text": content
                }
            ]
            if img_base is not None:
                content_list.append({
                    "type": "image_url",
                    "image_url": {
                        "url": img_base
                    }
                })
            msgs.append({'role':'user','content':content_list})
    response = client.chat.completions.create(
        model='glm-4v-flash',
        messages = msgs,
        temperature=0.6
    )
    return response.choices[0].message.content
if st.session_state.get('first_run',None) is None:
    st.write('第一次运行!')

    messages = [
    ('human','你好!'),
    ('ai','你好!我是智能聊天助理。可以帮你完成各种任务。'),
    ('human','你叫什么？')
]
    st.session_state['messages']=messages
    st.session_state['first_run'] = True

st.markdown('### 智能聊天机器人')

prompt = st.chat_input(
    "Say something and/or attach an image",
    accept_file=True,
    file_type=["jpg", "jpeg", "png"],
)


if prompt is not None:

    prompt_text = prompt['text']
    if len(prompt['files'])>0:
        img_base = base64.b64encode(prompt['files'][0].getvalue()).decode("utf-8")
        st.session_state['messages'].append(('human',prompt_text,img_base))
    else:
        st.session_state['messages'].append(('human',prompt_text))
    resp = llm_chat(st.session_state['messages'])
    st.session_state['messages'].append(('ai',resp))

for msg in st.session_state['messages']:
    with st.chat_message(msg[0]):
        st.write(msg[1])
